import os
import random
from dotenv import load_dotenv
import pymysql
import string

# 加载环境变量
load_dotenv(override=True)


# 预定义的网格ID和对应的经纬度
grid_mapping = {
    f"G{str(i).zfill(3)}": {
        "LONG": round(random.uniform(39.43, 41.05), 2),
        "LAT": round(random.uniform(115.25, 117.30), 2)
    } for i in range(1, 1000)
}


base_mapping = {
    f"B{str(i).zfill(3)}": {
        "LAT": round(random.uniform(39.43, 41.05), 2),
        "LONG": round(random.uniform(115.25, 117.30), 2)
    } for i in range(1, 1000)
}


# 生成随机字符串
def generate_random_string(prefix, length, start=1):
    return f"{prefix}{str(start + random.randint(0, int('9' * length))).zfill(length)}"


# 生成随机整数
def generate_random_int(min_val, max_val):
    return random.randint(min_val, max_val)


# 生成随机浮点数
def generate_random_float(min_val, max_val, precision=2):
    return round(random.uniform(min_val, max_val), precision)


# 生成省份编码
def generate_province():
    provinces = ["北京", "山东", "云南"]
    return random.choice(provinces)


# 生成统计时间
def generate_statis_time():
    return "20250731"


# 生成是否异网用户标签
def generate_is_diff_label():
    return random.choices([0, 1], weights=[70, 30])[0]

# 生成数据
def generate_data(num_records):
    data = []
    for i in range(num_records):
        is_diff = generate_is_diff_label()

        grid_id = generate_random_string("G", 3)
        grid_lat = grid_mapping.get(grid_id, {}).get("LAT")
        grid_long = grid_mapping.get(grid_id, {}).get("LONG")

        day_base_id = generate_random_string("B", 3) if is_diff == 0 else None
        day_base_lat = base_mapping.get(day_base_id, {}).get("LAT") if is_diff == 0 else None
        day_base_long = base_mapping.get(day_base_id, {}).get("LONG") if is_diff == 0 else None

        night_base_id = generate_random_string("B", 3) if is_diff == 0 else None
        night_base_lat = base_mapping.get(night_base_id, {}).get("LAT") if is_diff == 0 else None
        night_base_long = base_mapping.get(night_base_id, {}).get("LONG") if is_diff == 0 else None

        row = {
            "STATIS_TIME": generate_statis_time(),
            "PROV": generate_province(),
            "PERSON_ID": generate_random_string("P", 5),
            "USER_ID": generate_random_string("U", 5) if is_diff == 0 else None,
            "CUST_ID": generate_random_string("C", 5) if is_diff == 0 else None,
            "ATTR_AGE": generate_random_int(18, 60) if is_diff == 0 else None,
            "ATTR_GRID_ID": grid_id,
            "ATTR_GRID_LAT": grid_lat,
            "ATTR_GRID_LONG": grid_long,
            "ATTR_DAY_BASE_ID": day_base_id,
            "ATTR_DAY_BASE_LAT": day_base_lat,
            "ATTR_DAY_BASE_LONG": day_base_long,
            "ATTR_NIGHT_BASE_ID": night_base_id,
            "ATTR_NIGHT_BASE_LAT": night_base_lat,
            "ATTR_NIGHT_BASE_LONG": night_base_long,
            "LABEL_IS_GSM": generate_random_int(0, 1) if is_diff == 0 else None,
            "LABEL_IS_MID": generate_random_int(0, 1) if is_diff == 0 else None,
            "LABEL_IS_PREM": generate_random_int(0, 1) if is_diff == 0 else None,
            "LABEL_IS_DIFF": is_diff,
            "LABEL_IS_ELDER": generate_random_int(0, 1) if is_diff == 0 else None,
            "LABEL_IS_INT": generate_random_int(0, 1) if is_diff == 0 else None,
            "ATTR_FAMILY_CIRCLE_ID": generate_random_string("FC", 3),
            "ATTR_SOCIAL_CIRCLE_ID": generate_random_string("SC", 3),
            "ATTR_COMMUNITY_CIRCLE_ID": generate_random_string("CYC", 3),
            "ATTR_CAMPUS_CIRCLE_ID": generate_random_string("CPC", 3),
            "ATTR_COMMUTE_CIRCLE_ID": generate_random_string("CTC", 3),
            "ATTR_BUSINESS_CIRCLE_ID": generate_random_string("BC", 3),
            "ATTR_ONLINE_CIRCLE_ID": generate_random_string("OC", 3),
            "METRIC_ARPU": generate_random_float(100, 200) if is_diff == 0 else None,
            "METRIC_FLUX_USED": generate_random_float(1000000, 5000000) if is_diff == 0 else None,
            "METRIC_CALL_DURATION": generate_random_float(10, 1000) if is_diff == 0 else None
        }
        data.append(row)
    return data


# 连接到MySQL数据库
def connect_to_mysql(host, database, user, password):
    try:
        connection = pymysql.connect(
            host=host,
            database=database,
            user=user,
            password=password
        )
        print("Connected to MySQL database")
        return connection
    except pymysql.MySQLError as e:
        print(f"Error: {e}")
        return None


def check_and_create_table(connection, table_name):
    cursor = connection.cursor()

    # 检查表是否存在
    check_table_query = f"""
    SELECT COUNT(*)
    FROM information_schema.tables 
    WHERE table_schema = DATABASE() AND table_name = %s
    """
    cursor.execute(check_table_query, (table_name,))
    result = cursor.fetchone()[0]

    if result == 0:
        # 表不存在，创建表
        create_table_query = f"""
        CREATE TABLE {table_name} (
            STATIS_TIME VARCHAR(8),
            PROV VARCHAR(10),
            PERSON_ID VARCHAR(10),
            USER_ID VARCHAR(10),
            CUST_ID VARCHAR(10),
            ATTR_AGE INT,
            ATTR_GRID_ID VARCHAR(10),
            ATTR_GRID_LAT FLOAT,
            ATTR_GRID_LONG FLOAT,
            ATTR_DAY_BASE_ID VARCHAR(10),
            ATTR_DAY_BASE_LAT FLOAT,
            ATTR_DAY_BASE_LONG FLOAT,
            ATTR_NIGHT_BASE_ID VARCHAR(10),
            ATTR_NIGHT_BASE_LAT FLOAT,
            ATTR_NIGHT_BASE_LONG FLOAT,
            LABEL_IS_GSM TINYINT,
            LABEL_IS_MID TINYINT,
            LABEL_IS_PREM TINYINT,
            LABEL_IS_DIFF TINYINT,
            LABEL_IS_ELDER TINYINT,
            LABEL_IS_INT TINYINT,
            ATTR_FAMILY_CIRCLE_ID VARCHAR(10),
            ATTR_SOCIAL_CIRCLE_ID VARCHAR(10),
            ATTR_COMMUNITY_CIRCLE_ID VARCHAR(10),
            ATTR_CAMPUS_CIRCLE_ID VARCHAR(10),
            ATTR_COMMUTE_CIRCLE_ID VARCHAR(10),
            ATTR_BUSINESS_CIRCLE_ID VARCHAR(10),
            ATTR_ONLINE_CIRCLE_ID VARCHAR(10),
            METRIC_ARPU FLOAT,
            METRIC_FLUX_USED FLOAT,
            METRIC_CALL_DURATION FLOAT
        )
        """
        cursor.execute(create_table_query)
        connection.commit()
        print(f"Table {table_name} created.")
    else:
        print(f"Table {table_name} already exists.")


# 插入数据到MySQL表
def insert_data(connection, data, table_name):
    cursor = connection.cursor()

    insert_query = """
    INSERT INTO {} (
        STATIS_TIME, PROV, PERSON_ID, USER_ID, CUST_ID, 
        ATTR_AGE, ATTR_GRID_ID, ATTR_GRID_LAT, ATTR_GRID_LONG, 
        ATTR_DAY_BASE_ID, ATTR_DAY_BASE_LAT, ATTR_DAY_BASE_LONG, 
        ATTR_NIGHT_BASE_ID, ATTR_NIGHT_BASE_LAT, ATTR_NIGHT_BASE_LONG, 
        LABEL_IS_GSM, LABEL_IS_MID, LABEL_IS_PREM, LABEL_IS_DIFF, 
        LABEL_IS_ELDER, LABEL_IS_INT, 
        ATTR_FAMILY_CIRCLE_ID, ATTR_SOCIAL_CIRCLE_ID, ATTR_COMMUNITY_CIRCLE_ID, 
        ATTR_CAMPUS_CIRCLE_ID, ATTR_COMMUTE_CIRCLE_ID, ATTR_BUSINESS_CIRCLE_ID, 
        ATTR_ONLINE_CIRCLE_ID, METRIC_ARPU, METRIC_FLUX_USED, METRIC_CALL_DURATION
    ) VALUES (
        %s, %s, %s, %s, %s, 
        %s, %s, %s, %s, 
        %s, %s, %s, 
        %s, %s, %s, 
        %s, %s, %s, %s, 
        %s, %s, 
        %s, %s, %s, 
        %s, %s, %s, 
        %s, %s, %s, %s
    );
    """.format(table_name)
    for row in data:
        values = tuple(row.get(key, None) for key in [
            "STATIS_TIME", "PROV", "PERSON_ID", "USER_ID", "CUST_ID",
            "ATTR_AGE", "ATTR_GRID_ID", "ATTR_GRID_LAT", "ATTR_GRID_LONG",
            "ATTR_DAY_BASE_ID", "ATTR_DAY_BASE_LAT", "ATTR_DAY_BASE_LONG",
            "ATTR_NIGHT_BASE_ID", "ATTR_NIGHT_BASE_LAT", "ATTR_NIGHT_BASE_LONG",
            "LABEL_IS_GSM", "LABEL_IS_MID", "LABEL_IS_PREM", "LABEL_IS_DIFF",
            "LABEL_IS_ELDER", "LABEL_IS_INT",
            "ATTR_FAMILY_CIRCLE_ID", "ATTR_SOCIAL_CIRCLE_ID", "ATTR_COMMUNITY_CIRCLE_ID",
            "ATTR_CAMPUS_CIRCLE_ID", "ATTR_COMMUTE_CIRCLE_ID", "ATTR_BUSINESS_CIRCLE_ID",
            "ATTR_ONLINE_CIRCLE_ID", "METRIC_ARPU", "METRIC_FLUX_USED", "METRIC_CALL_DURATION"
        ])
        cursor.execute(insert_query, values)
    connection.commit()
    print(f"{len(data)} records inserted successfully")


def main():

    # 数据库连接信息
    host = os.getenv('HOST')
    user = os.getenv('USER')
    password = os.getenv('MYSQL_PW')
    database = os.getenv('DB_NAME')
    port = os.getenv('PORT')
    table_name = os.getenv('TABLE_NAME')

    connection = connect_to_mysql(host, database, user, password)
    check_and_create_table(connection, table_name)

    # 生成数据
    num_records = 200
    data = generate_data(num_records)

    # 插入数据
    insert_data(connection, data, table_name)

    # 关闭连接
    connection.close()


if __name__ == "__main__":
    main()

